import numpy as np
np.set_printoptions(threshold=np.inf)
import sys


def allclose_nparray(data_expected, data_me, rtol, atol, equal_nan=True):
    def _count_unequal_element(data_expected, data_me, rtol, atol):
        assert data_expected.shape == data_me.shape
        total_count = len(data_expected.flatten())
        print("total_count", total_count)
        error = np.abs(data_expected - data_me)/(np.abs(data_me) + np.min(np.abs(data_me))*1e-4) 
        greater = np.greater(error, atol + 1 * rtol)
        loss_count = np.count_nonzero(greater)
        print("loss_count", loss_count)
        assert (loss_count / total_count) < rtol, \
            "\ndata_expected_std:{0}\ndata_me_error:{1}\nloss:{2}, num: {3}/{4} {5}". \
                format(data_expected[greater], data_me[greater], error[greater], loss_count, total_count, loss_count/total_count)
                
    if np.any(np.isnan(data_expected)):
        assert np.allclose(data_expected, data_me, rtol, atol, equal_nan=equal_nan)
    elif not np.allclose(data_expected, data_me, rtol, atol, equal_nan=False):
        _count_unequal_element(data_expected, data_me, rtol, atol)
    else:
        assert True
        

def compare_func(compare1, compare2):
    try:
        precision = float(sys.argv[3])
    except:
        precision = float(1e-3)
    compare1_name  = compare1.split(".")[0]
    compare2_name  = compare2.split(".")[0]
    print("[INFO] Start test '{}' && '{}'!".format(compare1_name, compare2_name))
    compare1_data = np.load(compare1)
    compare2_data = np.load(compare2)
    
    print("[INFO] '{}' shape is {} && '{}' shape is {}!".format(compare1_name, compare1_data.shape, \
                                                            compare2_name, compare2_data.shape))
    compare1_data = compare1_data.reshape(compare2_data.shape)
    print(f"{compare1_name} data: {compare1_data.reshape(-1)[-100:]}")
    print(f"{compare2_name} data: {compare2_data.reshape(-1)[-100:]}")
    allclose_nparray(compare1_data, compare2_data, precision, 0)
    # try:
        # np.testing.assert_allclose(compare1_data, compare2_data, rtol=precision, atol=0)
        # print("[INFO] Test success with precision {}!".format(precision))
    # except Exception as ex:
        # print("[WARNING] Test failed with precision {}!".format(precision), end="")
        # print(ex)

if __name__ == "__main__":
    compare1 = sys.argv[1]
    compare2 = sys.argv[2]
    compare_func(compare1, compare2)

